import networkx as nx
import matplotlib.pyplot as plt
from itertools import islice


def build_hybrid_topology_shared_tor(num_tor):  # 参数化ToR数量
    """构建共享ToR的混合拓扑"""
    G = nx.Graph()

    # 设备前缀
    SERVER_PREFIX = "Server"
    TOR_PREFIX = "ToR"
    AGGR_ELECTRIC_PREFIX = "Aggr_E"
    CORE_ELECTRIC_PREFIX = "Core_E"
    OPTICAL_SWITCH_PREFIX = "Optical_SW"
    AGGR_OPTICAL_PREFIX = "Aggr_Optical"

    # ================== 共享ToR架构 ==================
    servers_per_tor = 1
    gpus_per_server = 8
    optical_ports_per_tor = num_tor  # 每个ToR的光端口数
    electric_ports_per_tor = 1  # 每个ToR的电端口数

    # 创建ToR交换机和服务器
    for tor_id in range(1, num_tor + 1):
        # 创建ToR节点（同时支持电和光）
        tor_node = f"ToR{tor_id}"  # 每个ToR交换机命名为ToR1
        G.add_node(tor_node, type='tor',
                   electric_capacity=40, # 每个ToR的电容量
                   optical_capacity=800, #每个ToR的光端口容量
                   electric_ports=electric_ports_per_tor, # 电端口数量
                   optical_ports=optical_ports_per_tor) # 光端口数量

        # 连接服务器
        for s in range(1, servers_per_tor + 1):
            server_node = f"{SERVER_PREFIX}{tor_id}-{s}" #定义服务器的命名规则
            G.add_node(server_node, type='server', gpus=gpus_per_server)
            G.add_edge(tor_node, server_node,
                       type='electrical', # 标记为电端口
                       capacity=400,  # 服务器到ToR的带宽
                       port_type='server')

    # ================== 电网络部分 ==================
    # 创建电网络核心层
    core_e_nodes = [f"{CORE_ELECTRIC_PREFIX}{i}" for i in range(1, int(num_tor/4.0+1))] # 定义核心交换机的命名规则
    G.add_nodes_from(core_e_nodes, type='core_electric', capacity=400) # 核心交换机的容量设为400(电网络）

    # 创建电网络聚合层
    aggr_e_nodes = [f"{AGGR_ELECTRIC_PREFIX}{i}" for i in range(1, int(num_tor/2.0+1))] #定义聚合交换机的命名规则
    G.add_nodes_from(aggr_e_nodes, type='aggregation_electric', capacity=100) # 聚合交换机的容量设为100

    # 连接电网络
    # 核心-聚合连接（每个核心连接2个聚合）
    for i, core in enumerate(core_e_nodes):
        for aggr_id in range(i * 2 + 1, i * 2 + 3):
            if aggr_id <= len(aggr_e_nodes):
                G.add_edge(core, f"{AGGR_ELECTRIC_PREFIX}{aggr_id}",
                           type='electrical', capacity=400)

    # 聚合-ToR连接（每个聚合连接2个ToR）
    for i, aggr in enumerate(aggr_e_nodes):
        for tor_id in range(i * 2 + 1, i * 2 + 3):
            if tor_id <= num_tor:
                G.add_edge(aggr, f"{TOR_PREFIX}{tor_id}",
                           type='electrical',
                           capacity=400,
                           port_type='uplink')  # 标记为电上行链路

    # 添加核心交换机之间的高速链路连接
    for i in range(len(core_e_nodes)):
        for j in range(i + 1, len(core_e_nodes)):
            G.add_edge(core_e_nodes[i], core_e_nodes[j],
                       type='electrical', capacity=1200)  # 高速链路容量设为1000

    # ================== 光网络部分 ==================
    # 创建光交换机
    optical_switches = [f"{OPTICAL_SWITCH_PREFIX}{i}" for i in range(1, num_tor+1)] # 创建光交换机的命名规则
    G.add_nodes_from(optical_switches, type='optical_switch', capacity=800) # 光交换机的容量设为1000

    # 连接光网络，让每个ToR连接所有光交换机
    for tor_id in range(1, num_tor + 1):
        tor_node = f"{TOR_PREFIX}{tor_id}"
        for sw_node in optical_switches:
            G.add_edge(tor_node, sw_node,
                       type='optical',
                       capacity=400,
                       port_type='optical')  # 标记为光端口

    return G


def visualize_topology(G):
    """可视化拓扑，分层排列节点"""
    # 定义节点颜色和形状
    node_colors = {
        'tor': 'lightblue',
        'server': 'lightgreen',
        'core_electric': 'orange',
        'aggregation_electric': 'yellow',
        'optical_switch': 'pink',
        # 'aggregation_optical': 'purple'
    }
    node_shapes = {
        'tor': 's',  # 方形
        'server': 'o',  # 圆形
        'core_electric': 'd',  # 菱形
        'aggregation_electric': '^',  # 三角形
        'optical_switch': 'p',  # 五边形
        # 'aggregation_optical': 'h'  # 六边形
    }

    # 定义边颜色
    edge_colors = {
        'electrical': 'gray',
        'optical': 'red'
    }

    # 手动指定节点位置
    pos = {}
    layer_spacing = 1.5  # 层间距
    node_spacing = 1.0  # 节点间距

    # 第一行：光交换机
    optical_switches = [n for n in G.nodes if G.nodes[n]['type'] == 'optical_switch']
    for i, node in enumerate(sorted(optical_switches)):
        pos[node] = (i * node_spacing, 4 * layer_spacing)

    # 第二行：ToR 交换机
    tor_switches = [n for n in G.nodes if G.nodes[n]['type'] == 'tor']
    for i, node in enumerate(sorted(tor_switches)):
        pos[node] = (i * node_spacing, 3 * layer_spacing)

    # 第三行：服务器
    servers = [n for n in G.nodes if G.nodes[n]['type'] == 'server']
    for i, node in enumerate(sorted(servers)):
        pos[node] = (i * node_spacing, 2 * layer_spacing)

    # 第四行：电聚合交换机
    aggr_electric = [n for n in G.nodes if G.nodes[n]['type'] == 'aggregation_electric']
    for i, node in enumerate(sorted(aggr_electric)):
        pos[node] = (i * node_spacing, 1 * layer_spacing)

    # 第五行：电核心交换机
    core_electric = [n for n in G.nodes if G.nodes[n]['type'] == 'core_electric']
    for i, node in enumerate(sorted(core_electric)):
        pos[node] = (i * node_spacing, 0 * layer_spacing)

    # 绘制节点
    plt.figure(figsize=(12, 8))
    for node_type, color in node_colors.items():
        nodes = [n for n in G.nodes if G.nodes[n]['type'] == node_type]
        nx.draw_networkx_nodes(G, pos, nodelist=nodes, node_color=color,
                               node_shape=node_shapes[node_type], label=node_type, node_size=300)

    # 绘制边
    for edge_type, color in edge_colors.items():
        edges = [(u, v) for u, v, d in G.edges(data=True) if d['type'] == edge_type]
        nx.draw_networkx_edges(G, pos, edgelist=edges, edge_color=color, label=edge_type, width=2)

    # 绘制标签
    nx.draw_networkx_labels(G, pos, font_size=8, font_weight='bold')

    # 显示图例
    plt.legend(scatterpoints=1, frameon=False, title="Node Types", bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.title("Hybrid Topology with Shared ToR (Layered Layout)", fontsize=14)
    plt.tight_layout()
    plt.show()


# 示例用法
if __name__ == "__main__":
    network = build_hybrid_topology_shared_tor(4)

    # 验证ToR连接
    tor1 = "ToR1"
    print(f"\n{tor1}的连接:")
    for neighbor in network.neighbors(tor1):
        edge_data = network.get_edge_data(tor1, neighbor)
        print(f" - {neighbor} ({edge_data['type']} link, port: {edge_data['port_type']})")

    # tor10 = "ToR10"
    # print(f"\n{tor10}的连接:")
    # for neighbor in network.neighbors(tor10):
    #     edge_data = network.get_edge_data(tor10, neighbor)
    #     print(f" - {neighbor} ({edge_data['type']} link, port: {edge_data['port_type']})")

    # 统计关键指标
    print("\n拓扑统计:")
    print(f"总节点数: {len(network.nodes)}")
    print(f"总连接数: {len(network.edges)}")
    print(f"电核心交换机: {len([n for n in network.nodes if 'core_electric' in network.nodes[n]['type']])}")
    print(f"光交换机: {len([n for n in network.nodes if 'optical_switch' in network.nodes[n]['type']])}")
    # print(f"光链路容量: {network.edges['core_electric1', 'optical_switch1']['capacity']}")
    # print(f"电链路容量: {sum([network.edges[u, v]['capacity'] for u, v in network.edges if 'electrical' in network.edges[u, v]['type']])}")
    # 输出每条链路的容量
    for u, v in network.edges():
        print(f"链路 {u}→{v} | 容量: {network.edges[u, v]['capacity']}")
    # 可视化拓扑
    visualize_topology(network)